# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import signal
import os

import tabulate
from tornado.ioloop import IOLoop

import coin.flow.subscriber as flow_subscriber
from coin.base.param_util import to_list
from coin.exchange.bitmex.kr_rest.currency import BitmexCurrency
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_query_pb2 import CurrencyBalance
from coin.strategy.mm.bitmex.order_subsystem import BitmexOrderSubsystem
from coin.strategy.mm.base.multi_feed_subsystem import MultiFeedSubsystem
from coin.strategy.mm.subscription import SubscriptionRequest


def calculate_unrealized_pnl_inverted(pos, avg_price, ask0_price, bid0_price):
  if pos == 0:
    return 0
  mid_price = (ask0_price + bid0_price) / 2.
  return pos * (1. / avg_price - 1. / mid_price)


def calculate_unrealized_pnl_ethusd(pos, avg_price, ask0_price, bid0_price):
  if pos == 0:
    return 0
  mid_price = (ask0_price + bid0_price) / 2.
  return pos * (mid_price - avg_price) * 0.000001


def calculate_unrealized_pnl(product, pos, avg_price, ask0_price, bid0_price):
  # TODO(jaewon): Performance
  if product == BitmexFuturesProduct.FromStr('ETH-USD.PERPETUAL'):
    return calculate_unrealized_pnl_ethusd(pos, avg_price, ask0_price, bid0_price)
  else:
    return calculate_unrealized_pnl_inverted(pos, avg_price, ask0_price, bid0_price)


class BitmexPositionBalancePrinter(object):
  def __init__(self, products, bookmap, currency=None, strategy_logger=None, logger=None):
    self._logger = logger or logging.getLogger('BitmexPositionBalancePrinter')
    self._products = sorted(to_list(products))
    self._bookmap = bookmap
    self._strategy_logger = strategy_logger

    currency = currency or 'BTC'
    if isinstance(currency, str):
      self._currency_str = currency
    else:
      self._currency_str = str(currency.currency)

    self._og = None
    self.started = False

  def register_order_gateway(self, order_gateway):
    self._og = order_gateway

  def start_print(self, ioloop, period_sec=10):
    if self.started:
      return
    self.started = True
    ioloop.add_timeout(datetime.timedelta(seconds=period_sec),
                       self._periodic_print,
                       ioloop,
                       period_sec)

  def _periodic_print(self, ioloop, period_sec):
    try:
      self.print_position_balance()
    finally:
      ioloop.add_timeout(datetime.timedelta(seconds=period_sec),
                         self._periodic_print,
                         ioloop,
                         period_sec)

  def _get_pos(self, product):
    pos = self._og.get_position(product)
    avg_entry_price = self._og.get_average_entry_price(product)

    book = self._bookmap.get_book(product)
    if book is None or not book.has_ask() or not book.has_bid():
      return None, None

    unrealized_pnl = calculate_unrealized_pnl(product,
                                              pos,
                                              avg_entry_price,
                                              book.ask0().price,
                                              book.bid0().price)
    return pos, unrealized_pnl

  def print_position_balance(self, ignore_error=False):
    try:
      if self._og is None or not self._og.is_ready():
        return

      table = []
      total = 0.
      for product in self._products:
        pos, unrealized_pnl = self._get_pos(product)
        total += (unrealized_pnl or 0.)
        table.append(['*', str(product), pos, unrealized_pnl])

      balance = self._og.get_balance(self._currency_str)
      total += balance

      table.append(['*', self._currency_str, None, balance])
      table.append(['*', 'Total', None, total])

      pos_bal_table_str = tabulate.tabulate(table,
                                            headers=['*', 'Product', 'Position', 'Value'],
                                            floatfmt='.4f')
      self._logger.info('----- Position and balance:\n* %s\n%s',
                        str(datetime.datetime.now()),
                        pos_bal_table_str)

      if self._strategy_logger:
        pnl_balance = CurrencyBalance(currency=self._currency_str, total=total)
        self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance, quote=None, subreq_id=None)
    except Exception:
      if ignore_error:
        self._logger.exception('Exception raised')
      else:
        raise

  def update_position_balance_status(self, status_out):
    if self._og is None or not self._og.is_ready():
      return

    upnl_sum = 0.
    for product in self._products:
      pos, unrealized_pnl = self._get_pos(product)
      if pos is None:
        continue

      status_out['position.%s' % str(product)] = '%.6f' % pos
      if unrealized_pnl is not None:
        upnl_sum += unrealized_pnl
        status_out['position.%s.unrealized_pnl' % str(product)] = ('%.6f' % unrealized_pnl)

    balance = self._og.get_balance(self._currency_str)
    status_out['balance.%s' % self._currency_str] = '%.6f' % balance
    status_out['balance.%s.total' % self._currency_str] = ('%.6f' % (balance + unrealized_pnl))


class BitmexStrategyBase(object):
  def __init__(self, products, logger=None):
    self._logger = logger or logging.getLogger('BitmexStrategyBase')
    self._products = to_list(products)
    self._ioloop = None
    self._queue = None
    self._flow_sub = None

    self._feed_subsystem = None
    self._og = None

  @property
  def order_gateway(self):
    return self._og

  @property
  def feed_subsystem(self):
    return self._feed_subsystem

  def register_flow(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._queue = queue
    self._flow_sub = flow_sub

  def register_feed_subsystem(self, feed_subsystem):
    self._feed_subsystem = feed_subsystem

  def register_order_gateway(self, name, order_gateway):
    self._logger.info('OG initialized: %s', name)
    if name == SubscriptionRequest('Futures', 'Bitmex', None):
      self._og = order_gateway
    else:
      raise ValueError('Unknown OG: %s' % name)

    self._ioloop.add_callback(self.on_og_reset)

  def prepare(self):
    pass

  def on_og_reset(self):
    raise NotImplementedError()

  def is_og_ready(self):
    return (self._og is not None and self._og.is_ready())

  def shutdown(self):
    self._ioloop.stop()


class SigtermHandler(object):
  def __init__(self, ioloop, callback):
    self._sig_count = 0
    self._ioloop = ioloop
    self._callback = callback

  def register_signal(self, sig):
    signal.signal(sig, self.handle_signal)

  def handle_signal(self, sig, frame):
    self._sig_count += 1
    try:
      logging.warning('Caught signal: sig=%s, count=%d', str(sig), self._sig_count)
      if self._sig_count == 1:
        self._ioloop.add_callback_from_signal(self._callback)
    finally:
      if self._sig_count == 3:
        os._exit(143)


def run_live(strategy,
             og_config_name,
             init_order_gateway_after=None,
             exit_after=None,
             install_sig_handler=True):
  # Initialize loop
  ioloop = IOLoop.current()
  queue = SimpleQueue('queue')
  flow_sub = flow_subscriber.from_queue(queue)
  strategy.register_flow(ioloop, queue, flow_sub)

  # Initialize feed
  feed_subsystem = MultiFeedSubsystem()
  strategy.register_feed_subsystem(feed_subsystem)

  # Initialize OG live environment
  env = BitmexOrderSubsystem(
      currencies=[BitmexCurrency.FromStr("BTC")],  # TODO(inkyu): Get it out
      trade_products=strategy._products,
      og_config_name=og_config_name,
      init_og_after=init_order_gateway_after,
      ioloop=ioloop,
      queue=queue,
      flow_sub=flow_sub,
      og_log_callback=strategy.strategy_logger.write_og_log)
  env.on_order_gateway_reset_callback.append(strategy.register_order_gateway)

  # Start
  strategy.prepare()
  feed_subsystem.start(ioloop, queue, flow_sub)
  env.start()

  if exit_after is not None:
    ioloop.add_timeout(exit_after, strategy.shutdown)

  if install_sig_handler:
    sigterm_handler = SigtermHandler(ioloop, strategy.shutdown)
    sigterm_handler.register_signal(signal.SIGTERM)
    sigterm_handler.register_signal(signal.SIGINT)

  try:
    ioloop.start()
    return 0
  except KeyboardInterrupt:
    return 1
